from Model import Actor, Distributions, Evidence, Model
from models.Crossingcrosswalk import CrossingCrosswalk
from models.Loiter import Loiter
from models.Run import Run
from models.Standstill import Standstill
from models.Walk import Walk
from segmentation.ChangepointFinder import ChangepointFinder
from segmentation.Formatter import Formatter
from segmentation.Move import MoveSegment
import argparse
import numpy as np
import os
import sys
import xml.etree.ElementTree as elet
    
DEBUG = 1
# models = [Loiter, Standstill, Walk, Run, Crossingcrosswalk]
models = [CrossingCrosswalk]

if __name__ == "__main__":
    
    # Parse command-line args (default to test values)
    parser = argparse.ArgumentParser(description='Run B3 Model Manager.')
    parser.add_argument('-t', '--tracks', default="./pol-regions.txt", #required=True,
                   help='Path to file containing the tracks evidence.')
    parser.add_argument('-o', '--outdir', default="./out/",  #required=True,
                   help='Path to output directory.')
    args = parser.parse_args()
    tracksfile = args.tracks
    outdir = args.outdir
    
    
    # Step 1: Read annotation into a list of actors
    my_formatter = Formatter()
    actor_list = my_formatter.process_pol_tracks_from_file(tracksfile)
    actor_list = sorted(actor_list, key=lambda a: int(a.name.split('_')[1]))
    
    
    # Step 2: Initialize a move-segmenter for each actor (using ChangepointFinder)
    actors_to_changepoint_finders = dict()
    actors_to_evidence = dict()
    names_to_actors = dict()
    for i,a in enumerate(actor_list):
        #if (a.name.split('_')[1] != '53'):
        if a.endtime - a.starttime > 300:
            continue
        if i >= 20:
            break
        names_to_actors[a.name] = a
        kw_dict = dict()
        arg_stat_funcs = [MoveSegment.identity_stat]
        arg_predictive_func = MoveSegment.expon_predictive
        arg_chi_prior = [.5]
        arg_nu_prior = 1.
        arg_actor = a
        arg_n_frames_for_params = 2
        arg_hazard_rate = .02
        arg_false_chi = [1./.15]
        arg_true_chi = [1./.8]
        arg_false_nu = 120.
        arg_true_nu = 60.
        arg_model_func = MoveSegment.expon_joint
        arg_model_params = [
            [Evidence.MOVE_FAST, [1./1.25]], [Evidence.MOVE_SLOW, [1./.75]],
            [Evidence.NOT_MOVE, [1./.05]]]
        # Create changepoint finder/segmenter for each actor a
        actors_to_changepoint_finders[a] = ChangepointFinder(
            param_func=MoveSegment.calculate_speed, stat_funcs=arg_stat_funcs,
            predictive_func=arg_predictive_func, chi_prior=arg_chi_prior,
            nu_prior=arg_nu_prior, actor=arg_actor,
            n_frames_for_params=arg_n_frames_for_params,
            hazard_rate=arg_hazard_rate, false_chi=arg_false_chi,
            false_nu=arg_false_nu, true_chi=arg_true_chi, true_nu=arg_true_nu,
            model_func=arg_model_func, model_params=arg_model_params)
        evidence_list = my_formatter.make_evidence_single(a)
        actors_to_evidence[a] = evidence_list
        test_actor = a
    
    
    # Step 3: Create TrackSegment evidence for each segmented interval.
    #  - Run changepoint segmenter for each actor. Output to XML.
    #  - Make TrackSegment evidence from XML output.
    aggregate_evidence_list = []
    for n, a in names_to_actors.iteritems():
        if DEBUG: print "test actor has properties %s" % a
        if DEBUG: print "test actor has width %s" % a.width
        evidence_list = actors_to_evidence[a]
        changepoint_finder = actors_to_changepoint_finders[a]
        n_frames = 0
        for e in evidence_list:
            n_frames += len(e.locations)
            if DEBUG: print "adding evidence with properties %s" % e
            changepoint_finder.add_segment(e)
        if DEBUG: print "added %s total frames" % n_frames
        
        # Generate changepoint for each actor a, output to XML.
        xml = changepoint_finder.generate_changepoint_xml()
        xml_tree = elet.ElementTree(xml)
        xml_root = xml_tree.getroot()
        xml_list = xml_root.findall("interval")
        
        # Make TrackSegment evidence for each interval
        segmented_evidence_list = a.make_evidence_from_xml(xml_list)
        
        # Aggregate all TrackSegments across all actors
        aggregate_evidence_list.extend(segmented_evidence_list)
        
    if DEBUG: print "ALL EVIDENCE: " + " ".join([str(x) for x in aggregate_evidence_list])
    
    
    # Step 4: Instantiate relevant activity models for each piece of evidence
    for ModelClass in models:
        test_model = ModelClass()
        for e in aggregate_evidence_list:
            test_model.model.add_evidence(e)
        test_model.model.evaluate_evidence()
        test_model.model.name
        test_model.model.print_xml()
#        test_model.model.print_xml(outf = os.path.join(outdir, test_model.model.name), mfiles=True)
